/*
 * @Author: LVGRAPE
 * @LastEditors: LVGRAPE
 */
#include "estimator.h"
#include "fc.h"
#include <math.h>
#include "anoPTv8_report.h"
#include "zino.h"
#include "position.h"
#include "filters.h"
#include "pwm3901_polling.h"

#define DBG_TAG "EST"
#define DBG_LVL DBG_LOG
#include <rtdbg.h>

#define GRAVITY_CMSS (980.665f) // 单位：cm/s²
#define GRAVITY_MSS (9.80665f)  // 单位：m/s²
#define ACC_DEADBAND 20
estimator_t fcEstimator;

float baroAlt, baroAlt_pre, baro_vel;
float tofAlt, tofAlt_pre, tof_vel;
uint16_t tofRange_pre;
float tofAlt_fix; // avoid tofAlt jump
bool isTofAltValid = false;

float x_est[2] = {0.0f, 0.0f};              // 位置，速度
float y_est[2] = {0.0f, 0.0f};              // 位置，速度
float z_est[2] = {0.0f, 0.0f};              // 位置，速度
float acc[3] = {0.0f, 0.0f, 0.0f};          // 加速度计数据
float acc_lpf[3] = {0.0f, 0.0f, 0.0f};      // 加速度计低通滤波数据
float accel_bias_corr[3] = {0.0, 0.0, 0.0}; // 加速度计偏置校正
float obs_height_lpf;                       // cm
float obs_height;                           // cm

float params_w_baro_p = 2;
float params_w_baro_v = 3;
float params_w_tof_p = 0.5;
float params_w_tof_v = 1;
float params_w_opt_p = 1;
float params_w_opt_v = 1;
float corr_baro;
float corr_baro_vel;
float corr_tof;
float corr_tof_vel;
float corr_opt_vx;
float corr_opt_vy;
float corr_opt_px;
float corr_opt_py;

bool params_reload_flag = false;

float baro_auto_gain = 1.0f;
int16_t opt_quality_lpf;
pt1Filter_t accZ_fiter;
pt1Filter_t baro_v_fiter;

const struct emstimator_params est_params_default = {
    .saved_flag = 0,
    .w_baro_p = 5,
    .w_baro_v = 9,
    .w_baro_a = 1,
    .w_tof_p = 9,
    .w_tof_v = 0,
    .w_tof_a = 0,
    .w_opt_p = 0,
    .w_opt_v = 0,
    .w_opt_a = 0,
};

struct emstimator_params est_params;
struct emstimator_params_f est_params_f;

ZINO_CONFIG_REGISTER(struct emstimator_params, estimator_params_config, 0, 1);

float get_tof_alt(void)
{
    return tofAlt;
}
void estimator_init(void)
{
    if (estimator_params_config()->saved_flag == ESTIMATOR_SAVED_FLAG)
    {
        rt_memcpy(&est_params, estimator_params_config(), sizeof(struct emstimator_params));
    }
    else
    {
        rt_memcpy(&est_params, &est_params_default, sizeof(struct emstimator_params));
        rt_memcpy(estimator_params_config(), &est_params_default, sizeof(struct emstimator_params));
    }

    est_params_f.w_baro_p = est_params.w_baro_p / 10.0f;
    est_params_f.w_baro_v = est_params.w_baro_v / 10.0f;
    est_params_f.w_baro_a = est_params.w_baro_a / 10.0f;
    est_params_f.w_tof_p = est_params.w_tof_p / 10.0f;
    est_params_f.w_tof_v = est_params.w_tof_v / 10.0f;
    est_params_f.w_tof_a = est_params.w_tof_a / 10.0f;
    est_params_f.w_opt_p = est_params.w_opt_p / 10.0f;
    est_params_f.w_opt_v = est_params.w_opt_v / 10.0f;
    est_params_f.w_opt_a = est_params.w_opt_a / 10.0f;

    pt1FilterInit(&accZ_fiter, pt1FilterGain(1, 0.01f));
    pt1FilterInit(&baro_v_fiter, pt1FilterGain(20, 0.01f));
}
void reload_estimator_params(void)
{
    if (params_reload_flag)
    {
        if (estimator_params_config()->saved_flag == ESTIMATOR_SAVED_FLAG)
        {
            rt_memcpy(&est_params, estimator_params_config(), sizeof(struct emstimator_params));
            params_reload_flag = false;
            z_est[ALT_POS] = baroAlt;
            z_est[ALT_VEL] = 0.0f;
            accel_bias_corr[Z] = 0;

            est_params_f.w_baro_p = est_params.w_baro_p / 10.0f;
            est_params_f.w_baro_v = est_params.w_baro_v / 10.0f;
            est_params_f.w_baro_a = est_params.w_baro_a / 10.0f;
            est_params_f.w_tof_p = est_params.w_tof_p / 10.0f;
            est_params_f.w_tof_v = est_params.w_tof_v / 10.0f;
            est_params_f.w_tof_a = est_params.w_tof_a / 10.0f;
            est_params_f.w_opt_p = est_params.w_opt_p / 10.0f;
            est_params_f.w_opt_v = est_params.w_opt_v / 10.0f;
            est_params_f.w_opt_a = est_params.w_opt_a / 10.0f;

            LOG_D("estimator params reload success");
            LOG_D("w_baro: %3d, %3d, %3d", est_params.w_baro_p, est_params.w_baro_v, est_params.w_baro_a);
            LOG_D("w_tof: %3d, %3d, %3d", est_params.w_tof_p, est_params.w_tof_v, est_params.w_tof_a);
            LOG_D("w_opt: %3d, %3d, %3d", est_params.w_opt_p, est_params.w_opt_v, est_params.w_opt_a);
        }
    }
}
void save_estimator_params(void)
{
    estimator_params_config()->saved_flag = ESTIMATOR_SAVED_FLAG;
    params_reload_flag = true;
}
/*
 * 函数名：inertial_filter_predict
 * 功能：一次预测
 */
void inertial_filter_predict(float dt, float x[2], float acc)
{
    if (isfinite(dt))
    {
        if (!isfinite(acc))
        {
            acc = 0.0f;
        }
        x[0] += x[1] * dt + acc * dt * dt / 2.0f;
        x[1] += acc * dt;
    }
}

/*
 * 函数名：inertial_filter_correct
 * 功能：一次校正
 */
void inertial_filter_correct(float e, float dt, float x[2], int i, float w)
{
    if (isfinite(e) && isfinite(w) && isfinite(dt))
    {
        float ewdt = e * w * dt;
        x[i] += ewdt;
        // if (i == 0)
        // {
        // x[1] += w * w * ewdt;
        // }
    }
}
void estimator_xy_update(float dt)
{
    if (dt < 0.01f)
    {
        dt = 0.01f; // 防止dt过小
    }

    fc_optical.gfix_x = constrainf(fc_optical.gfix_x, -20, 20);//低于30cm时，光流数据变强很多
    fc_optical.gfix_y = constrainf(fc_optical.gfix_y, -20, 20);//低于30cm时，光流数据变强很多
    fc_optical.hfix_x = pmw3901mb_raw_to_cm(fc_optical.gfix_x, obs_height_lpf * 0.01f) / dt;
    fc_optical.hfix_y = pmw3901mb_raw_to_cm(fc_optical.gfix_y, obs_height_lpf * 0.01f) / dt;

    acc[X] = (zdrone_state.acc.x + accel_bias_corr[X]) * GRAVITY_CMSS;
    acc[Y] = (zdrone_state.acc.y + accel_bias_corr[Y]) * GRAVITY_CMSS;

    x_est[HOR_VEL] += acc[X] * dt;
    y_est[HOR_VEL] += acc[Y] * dt;

    corr_opt_vx = fc_optical.hfix_x - x_est[HOR_VEL];
    corr_opt_vy = fc_optical.hfix_y - y_est[HOR_VEL];

    inertial_filter_correct(corr_opt_vx, dt, x_est, HOR_VEL, est_params_f.w_opt_v);
    inertial_filter_correct(corr_opt_vy, dt, y_est, HOR_VEL, est_params_f.w_opt_v);

    fc_optical.afix_x = applyDeadbandFloat(x_est[HOR_VEL], 1.f);
    fc_optical.afix_y = applyDeadbandFloat(y_est[HOR_VEL], 1.f);

    x_est[HOR_POS] = x_est[HOR_VEL] * dt + 0.5f * acc[X] * dt * dt;
    y_est[HOR_POS] = y_est[HOR_VEL] * dt + 0.5f * acc[Y] * dt * dt;

    fc_optical.pos_x = x_est[HOR_POS];
    fc_optical.pos_y = y_est[HOR_POS];

    opt_quality_lpf += (fc_optical.quality - opt_quality_lpf) >> 2;
    if (opt_quality_lpf < 90)
    {
        x_est[HOR_VEL] = 0;
        y_est[HOR_VEL] = 0;
    }

    fc_optical.is_valid = 1;
    fc_optical.is_withHeight = 1;
    fc_optical.is_withIMU = 1;

    zdrone_state.vel.x = fc_optical.afix_x;
    zdrone_state.vel.y = fc_optical.afix_y;
    zdrone_state.pos.x = fc_optical.pos_x;
    zdrone_state.pos.y = fc_optical.pos_y;

    if (FC_STATE_CHECK(FC_IMU_CALIBRATING))
    {
        x_est[HOR_POS] = 0;
        y_est[HOR_POS] = 0;
        x_est[HOR_VEL] = 0;
        y_est[HOR_VEL] = 0;
    }
}
void estimator_update(float dt)
{
    if (dt < 0.01f)
        dt = 0.01f; // 防止dt过小
    if (!FC_STATE_CHECK(FC_ARMED))
    {
        // 未解锁时，加大baro权重
        baro_auto_gain = 10.f;
    }
    else
    {
        if (!FC_STATE_CHECK(FC_AIRBORN)) // 起飞时权重小，避免地面效应
        {
            baro_auto_gain = 0.1f;
        }
        else
        {
            if (baro_auto_gain < 1.0f)
            {
                baro_auto_gain += 0.01f;
            }
            else
            {
                baro_auto_gain = 1.0f;
            }
        }
    }

    if (FC_STATE_CHECK(FC_IMU_CALIBRATING) || (millis() < 500))
    {
        fc_sensor_data.baro.offset = fc_sensor_data.baro.altitude;
        tofAlt = 0;
        accel_bias_corr[Z] = 0;
        z_est[ALT_POS] = 0;
        z_est[ALT_VEL] = 0;
    }
    baroAlt = fc_sensor_data.baro.altitude - fc_sensor_data.baro.offset; // 单位：cm
    baro_vel = pt1FilterApply(&baro_v_fiter, (baroAlt - baroAlt_pre) / dt);

    baro_vel = constrainf(baro_vel, -200.0f, 200.0f);
    baroAlt_pre = baroAlt;
    acc[Z] = (zdrone_state.acc.z + accel_bias_corr[Z]) * GRAVITY_CMSS;
    acc_lpf[Z] = pt1FilterApply(&accZ_fiter, acc[Z]);

    if (acc_lpf[Z] > ACC_DEADBAND || acc_lpf[Z] < -ACC_DEADBAND)
    {
        // boost acc vel weight
    }

    inertial_filter_predict(dt, z_est, acc[Z]);
    // z_est[ALT_VEL] = applyDeadbandFloat(z_est[ALT_VEL], 10.5f);
    corr_baro = baroAlt - z_est[ALT_POS];
    corr_baro_vel = baro_vel - z_est[ALT_VEL];

    inertial_filter_correct(corr_baro, dt, z_est, ALT_POS, est_params_f.w_baro_p * baro_auto_gain);
    inertial_filter_correct(corr_baro_vel, dt, z_est, ALT_VEL, est_params_f.w_baro_v * baro_auto_gain);
    accel_bias_corr[Z] += corr_baro_vel * est_params_f.w_baro_a * dt * dt * dt;

    isTofAltValid = fc_sensor_data.tof.range < 6000 && fc_sensor_data.tof.quality > 0 && fc_sensor_data.tof.isOnline;
    if (isTofAltValid)
    {
        int16_t dtTof = fc_sensor_data.tof.range - tofRange_pre;
        tofRange_pre = fc_sensor_data.tof.range;
        if (ABS(dtTof) < 50)
        {
            tofAlt += (dtTof * 0.1f); // cm
        }
        corr_tof = tofAlt - z_est[ALT_POS];
        tof_vel = (tofAlt - tofAlt_pre) / dt;
        tofAlt_pre = tofAlt;
        corr_tof_vel = tof_vel - z_est[ALT_VEL];
        inertial_filter_correct(corr_tof, dt, z_est, ALT_POS, est_params_f.w_tof_p * USED_TOF);
        inertial_filter_correct(corr_tof_vel, dt, z_est, ALT_VEL, est_params_f.w_tof_v * USED_TOF);
        accel_bias_corr[Z] += corr_tof_vel * est_params_f.w_tof_a * dt * dt * dt;

        obs_height = fc_sensor_data.tof.range * 0.1f;
    }
    else
    {
        tofAlt = z_est[ALT_POS];
        obs_height = z_est[ALT_POS];
    }
    obs_height = constrainf(obs_height, 1, 500);
    obs_height_lpf += (obs_height - obs_height_lpf) * 0.1f;

    estimator_xy_update(dt);

    z_est[ALT_VEL] = constrainf(z_est[ALT_VEL], -200.0f, 200.0f);
    zdrone_state.pos.z = z_est[ALT_POS];
    zdrone_state.vel.z = applyDeadbandFloat(z_est[ALT_VEL], 1);
    zdrone_state.pos.x = x_est[HOR_POS];
    zdrone_state.vel.x = applyDeadbandFloat(x_est[HOR_VEL], 1);
    zdrone_state.pos.y = y_est[HOR_POS];
    zdrone_state.vel.y = applyDeadbandFloat(y_est[HOR_VEL], 1);

    F1_wave[0] = fc_optical.raw_x;
    F1_wave[1] = fc_optical.raw_y;
    F1_wave[2] = fc_optical.gfix_x;
    F1_wave[3] = fc_optical.gfix_y;
    F1_wave[4] = fc_optical.hfix_x;
    F1_wave[5] = fc_optical.hfix_y;
    F1_wave[6] = fc_optical.afix_x;
    F1_wave[7] = fc_optical.afix_y;
    F1_wave[8] = fc_optical.is_valid;
    F1_wave[9] = get_tof_alt();
    F1_wave[10] = zdrone_state.vel.z;
    F1_wave[11] = zdrone_state.pos.z;

    // F1_wave[0] = baroAlt;
    // F1_wave[1] = baro_vel;
    // F1_wave[2] = tofAlt;
    // F1_wave[3] = tof_vel;
    // F1_wave[4] = z_est[ALT_POS];
    // F1_wave[5] = z_est[ALT_VEL];
}
void estimator_reset_params(void)
{
    rt_memcpy(estimator_params_config(), &est_params_default, sizeof(struct emstimator_params));
}
